"""
实时股价数据API
提供WebSocket和REST接口
"""

import json
import logging
from typing import List, Optional
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException, Query
from fastapi.responses import JSONResponse

from ..services.realtime_service import realtime_service

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/realtime", tags=["实时数据"])

@router.on_event("startup")
async def startup_realtime_service():
    """启动时初始化实时服务"""
    await realtime_service.init_redis()
    # 启动实时更新任务（在后台运行）
    import asyncio
    asyncio.create_task(realtime_service.start_realtime_updates())

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket连接端点"""
    await realtime_service.connect_websocket(websocket)
    
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            message = json.loads(data)
            
            action = message.get("action")
            stock_code = message.get("stock_code")
            
            if action == "subscribe" and stock_code:
                await realtime_service.subscribe_stock(stock_code)
                await websocket.send_text(json.dumps({
                    "type": "subscription_success",
                    "stock_code": stock_code,
                    "message": f"已订阅 {stock_code} 实时数据"
                }, ensure_ascii=False))
                
            elif action == "unsubscribe" and stock_code:
                await realtime_service.unsubscribe_stock(stock_code)
                await websocket.send_text(json.dumps({
                    "type": "unsubscription_success", 
                    "stock_code": stock_code,
                    "message": f"已取消订阅 {stock_code}"
                }, ensure_ascii=False))
                
            elif action == "ping":
                await websocket.send_text(json.dumps({
                    "type": "pong",
                    "timestamp": message.get("timestamp")
                }))
                
    except WebSocketDisconnect:
        await realtime_service.disconnect_websocket(websocket)
    except Exception as e:
        logger.error(f"WebSocket错误: {e}")
        await realtime_service.disconnect_websocket(websocket)

@router.get("/price/{stock_code}")
async def get_realtime_price(stock_code: str):
    """获取单个股票实时价格"""
    try:
        data = await realtime_service.get_realtime_price(stock_code)
        return JSONResponse(content={
            "success": True,
            "data": data,
            "message": "获取实时价格成功"
        })
    except Exception as e:
        logger.error(f"获取实时价格失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取实时价格失败: {str(e)}")

@router.get("/prices")
async def get_multiple_prices(
    stock_codes: str = Query(..., description="股票代码，用逗号分隔")
):
    """获取多个股票实时价格"""
    try:
        codes = [code.strip() for code in stock_codes.split(",")]
        
        if len(codes) > 20:
            raise HTTPException(status_code=400, detail="最多支持20个股票代码")
        
        import asyncio
        tasks = [realtime_service.get_realtime_price(code) for code in codes]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        data = []
        for i, result in enumerate(results):
            if isinstance(result, dict):
                data.append(result)
            else:
                logger.error(f"获取 {codes[i]} 数据失败: {result}")
        
        return JSONResponse(content={
            "success": True,
            "data": data,
            "count": len(data),
            "message": f"获取 {len(data)} 个股票实时价格成功"
        })
        
    except Exception as e:
        logger.error(f"获取多个实时价格失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取实时价格失败: {str(e)}")

@router.get("/minute/{stock_code}")
async def get_minute_data(
    stock_code: str,
    period: str = Query("1d", description="时间周期: 1d, 5d, 1w")
):
    """获取分钟级历史数据"""
    try:
        data = await realtime_service.get_stock_history_minute(stock_code, period)
        return JSONResponse(content={
            "success": True,
            "data": data,
            "count": len(data),
            "stock_code": stock_code,
            "period": period,
            "message": "获取分钟数据成功"
        })
    except Exception as e:
        logger.error(f"获取分钟数据失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取分钟数据失败: {str(e)}")

@router.get("/status")
async def get_realtime_status():
    """获取实时服务状态"""
    try:
        redis_status = "connected" if realtime_service.redis_client else "disconnected"
        
        return JSONResponse(content={
            "success": True,
            "data": {
                "service_running": realtime_service.is_running,
                "active_connections": len(realtime_service.active_connections),
                "subscribed_stocks": list(realtime_service.subscribed_stocks),
                "redis_status": redis_status
            },
            "message": "获取服务状态成功"
        })
    except Exception as e:
        logger.error(f"获取服务状态失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取服务状态失败: {str(e)}")

@router.post("/subscribe/{stock_code}")
async def subscribe_stock_rest(stock_code: str):
    """REST方式订阅股票（用于测试）"""
    try:
        await realtime_service.subscribe_stock(stock_code)
        return JSONResponse(content={
            "success": True,
            "stock_code": stock_code,
            "message": f"已订阅 {stock_code} 实时数据"
        })
    except Exception as e:
        logger.error(f"订阅股票失败: {e}")
        raise HTTPException(status_code=500, detail=f"订阅失败: {str(e)}")

@router.delete("/subscribe/{stock_code}")
async def unsubscribe_stock_rest(stock_code: str):
    """REST方式取消订阅股票"""
    try:
        await realtime_service.unsubscribe_stock(stock_code)
        return JSONResponse(content={
            "success": True,
            "stock_code": stock_code,
            "message": f"已取消订阅 {stock_code}"
        })
    except Exception as e:
        logger.error(f"取消订阅失败: {e}")
        raise HTTPException(status_code=500, detail=f"取消订阅失败: {str(e)}")

@router.get("/test")
async def test_realtime_api():
    """测试实时API"""
    return JSONResponse(content={
        "success": True,
        "message": "实时数据API正常运行",
        "endpoints": {
            "websocket": "/api/realtime/ws",
            "single_price": "/api/realtime/price/{stock_code}",
            "multiple_prices": "/api/realtime/prices?stock_codes=000001,000002",
            "minute_data": "/api/realtime/minute/{stock_code}",
            "status": "/api/realtime/status"
        }
    })
